from .hook import AfterIterHook
from ...metric.base import Metric


class MetricHook(AfterIterHook):

    def __init__(self, metric: Metric, name: str = "MetricHook"):
        super().__init__(name)
        self.metric = metric

    def after_iter(self, container):
        pred = container.pred
        _, label = container.data
        self.metric.update(pred, label)
        name, metric = self.metric.get()
        self.log("%s %.4f" % (name, metric), container)
